from pydantic import BaseModel, Field
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document

# Prompt工程——提升回答质量, 强制结构化输出, 引用与来源

from langchain_ollama import ChatOllama
llm = ChatOllama(model="qwen3:8b", temperature=0.5, reasoning=False)
from langchain_ollama import OllamaEmbeddings
embeddings = OllamaEmbeddings(model="bge-m3:latest")

documents = [    
    Document(page_content="LCEL，全称LangChain Expression Language，它通过操作符重载（如`|`符号）提供了一种声明式的、流畅的方式来构建AI链。它的关键优势包括：开箱即用的流式处理、异步和并行执行能力，以及对整个链的生命周期管理（如日志、调试）提供了极大的便利。", metadata={"doc_id": "lcel_intro"}),
    Document(page_content="混合搜索（Hybrid Search）结合了传统关键词搜索（如BM25）和现代向量搜索的优点。关键词搜索能精确匹配术语和缩写，而向量搜索擅长理解语义和意图。二者结合能显著提升检索的鲁棒性和准确性。", metadata={"doc_id": "hybrid_search_intro"}),
]

vectorstore = FAISS.from_documents(documents, embeddings)

class AnswerFormat(BaseModel):    
    answer: str = Field(description="针对问题的回答")    
    references: list[str] = Field(description="答案引用的文档块ID或来源信息列表")    
    confidence_score: float = Field(description="对答案的自信度评分,0.0到1.0之间")

# 绑定 LLM, 强制其输出 AnswerFormat 结构
llm_structured_output = llm.with_structured_output(AnswerFormat)

# 定义一个 Prompt,鼓励LLM输出结构化数据
structured_rag_prompt = ChatPromptTemplate.from_messages(
    [    
    ("system", "你是一个知识问答机器人。请根据提供的上下文,以JSON格式回答问题。如果信息不足,请将answer字段设置为'不知道',confidence_score为0.0。"),    
    ("user", "问题:{question}\n\n上下文:\n{context}"),
    ]
)

# 创建一个只负责生成结构化答案的链
structured_answer_chain = (    
    {"question": RunnablePassthrough(), "context": vectorstore.as_retriever()}    
    | structured_rag_prompt    
    | llm_structured_output)

print("\n--- 结构化输出 RAG 示例 ---")
query_structured = "LangChain的优势是什么？?"
response_structured = structured_answer_chain.invoke(query_structured)

print(f"问题: {query_structured}")
# pydantic模型可以方便地转为JSON
print(f"回答:\n{response_structured}")
